import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OneHotEncoder, StandardScaler


def preprocess_categorical(df, categorical_columns):
    """分类特征预处理"""
    # 创建OneHot编码器
    ohe = OneHotEncoder(handle_unknown="ignore", sparse_output=False)

    # 创建列转换器
    ct = ColumnTransformer([("cat", ohe, categorical_columns)], remainder="passthrough")

    # 应用预处理
    X_processed = ct.fit_transform(df)

    return X_processed, ct


if __name__ == "__main__":
    # 示例数据
    df = pd.DataFrame(
        {
            "court": ["A", "B", "A", "C"],
            "field": ["X", "Y", "Z", "Y"],
            "score": [80, 90, 75, 85],
            "attendance": [100, 150, 80, 120],
        }
    )

    # 打印原始数据
    print("原始数据:")
    print(df)
    print("\n" + "=" * 50 + "\n")

    # 定义分类特征列
    categorical_columns = ["court", "field"]

    # 应用预处理
    X_processed, preprocessor = preprocess_categorical(df, categorical_columns)

    # 打印预处理后的数据
    print("预处理后的数据 (NumPy数组):")
    print(X_processed)
    print(f"形状: {X_processed.shape}")
    print("\n" + "=" * 50 + "\n")

    # 获取特征名称
    feature_names = []
    for name, transformer, columns in preprocessor.transformers_:
        if name == "cat":
            # 获取OneHot编码后的特征名称
            ohe = transformer
            feature_names.extend(ohe.get_feature_names_out(columns))
        else:
            # 保留原始数值特征名称
            feature_names.extend(columns)

    # 转换为DataFrame以便更好地查看
    df_processed = pd.DataFrame(X_processed, columns=feature_names)
    print("预处理后的数据 (DataFrame格式):")
    print(df_processed)
    print("\n" + "=" * 50 + "\n")

    # 测试新数据的预处理
    print("测试新数据预处理:")
    new_data = pd.DataFrame(
        {
            "court": ["B", "D"],  # D是未见过的类别
            "field": ["Y", "X"],
            "score": [95, 70],
            "attendance": [130, 90],
        }
    )

    print("新数据:")
    print(new_data)

    # 使用训练好的预处理器处理新数据
    new_processed = preprocessor.transform(new_data)
    df_new_processed = pd.DataFrame(new_processed, columns=feature_names)
    print("\n新数据预处理结果:")
    print(df_new_processed)
